#全连接神经网络可视化
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
from tensorflow.contrib.tensorboard.plugins import projector
import os

input_size = 784
class_num = 10
#训练次数
step_num = 1001
#图片数量
img_num = 3000

#数据集
mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)

#载入图片
embedding = tf.Variable(tf.stack(mnist.test.images[:img_num]), name='embedding', trainable=False)

# 可视化标签
with tf.name_scope('input'):
    x = tf.placeholder(tf.float32, [None, input_size],name='x-input')
    y = tf.placeholder(tf.float32, [None, class_num],name='y-input')


def v_summaries(var):
    with tf.name_scope('v_summaries'):
        mean = tf.reduce_mean(var)
        tf.summary.scalar('mean', mean)
        stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean)))
        tf.summary.scalar('stddev', stddev)
        tf.summary.histogram('histogram', var)

# 构建一层
with tf.name_scope('layer'):
    with tf.name_scope('weight'):
        w = tf.Variable(tf.zeros([784, 10]), name='w')
        v_summaries(w)
    with tf.name_scope('biases'):
        b = tf.Variable(tf.zeros([10]), name='b')
        v_summaries(b)
    with tf.name_scope('wx_add_b'):
        wx_add_b = tf.add(tf.matmul(x,w),b)
    with tf.name_scope('softmax'):
        pred = tf.nn.softmax(wx_add_b)

with tf.name_scope('loss'):
    loss = tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=pred)
    #tf.summary.scalar('loss', loss)

with tf.name_scope('train'):
    train = tf.train.AdamOptimizer(0.1).minimize(loss)

with tf.name_scope('accr'):
    corr = tf.equal(tf.argmax(y,1),tf.argmax(pred,1))
    accr = tf.reduce_mean(tf.cast(corr, tf.float32))
    tf.summary.scalar('accr', accr)

merged = tf.summary.merge_all()
init = tf.global_variables_initializer()

# 生成meta data 文件
dir = "D:/py_workspace/py_example/8_tensor/projector"
if os.path.exists(dir+'/metadata.tsv'):
    os.remove(dir+'/metadata.tsv')

with tf.Session() as sess:
    sess.run(init)
    with open(dir + '/metadata.tsv', 'w') as f:
        labels = sess.run(tf.argmax(mnist.test.labels[:],1))
        for i in range(img_num):
            f.write(str(labels[i])+'\n')
    projector_writer = tf.summary.FileWriter(dir,sess.graph)
    saver = tf.train.Saver()
    config = projector.ProjectorConfig()
    embed = config.embeddings.add()
    embed.tensor_name = embedding.name
    embed.metadata_path = dir + '/metadata.tsv'
    embed.sprite.image_path = dir + '/imgs/mnist_10k_sprite.png'
    embed.sprite.single_image_dim.extend([28,28])
    projector.visualize_embeddings(projector_writer, config)

    for i in range(step_num):
        batch_xs, batch_ys = mnist.train.next_batch(100)
        run_options = tf.RunOptions()
        run_metadata = tf.RunMetadata()
        summary,_ = sess.run([merged, train],feed_dict={x:batch_xs,y:batch_ys},options=run_options,run_metadata=run_metadata)
        projector_writer.add_run_metadata(run_metadata, 'step %d' % i)
        projector_writer.add_summary(summary,i)
        if i%100 == 0:
            acc = sess.run(accr,feed_dict={x:mnist.test.images,y:mnist.test.labels})
            print('setp:'+str(i)+', test accr:'+str(acc))
        if i == (step_num-1):
            saver.save(sess,dir + '/model.ckpt',global_step=step_num)
projector_writer.close()
